/**
 * @file rls.h
 * @author Star_Plucking
 * @brief 递归最小二乘法
 * @version 0.1
 * @date 2025-02-26
 *
 */

#pragma once

#include <algorithm/math.h>
#include <algorithm/matrix.h>
#include <cstdint>

namespace os::math {
template <uint32_t Dim>
class RLS {
private:
    float lambda_;
    float delta_;
    uint32_t dimension_; ///> 维数
    Matrix<Dim, Dim> transMatrix_; ///> 协方差矩阵
    Matrix<Dim, 1> gainVector_; ///> 增益向量
    Matrix<Dim, 1> paramsVector_; ///> 参数向量
    Matrix<Dim, 1> defaultParamsVector_; ///> 初始参数向量

public:
    RLS() = default;
    ~RLS() = default;

    constexpr RLS(float lambda, float delta) : lambda_(lambda), delta_(delta), dimension_(Dim) {
        transMatrix_ = Matrix<Dim, Dim>::eye() * delta_;
        gainVector_ = Matrix<Dim, 1>::zeros();
        paramsVector_ = Matrix<Dim, 1>::zeros();
        defaultParamsVector_ = Matrix<Dim, 1>::zeros();
    }

    constexpr RLS(float delta_, float lambda_, Matrix<Dim, 1> initParam) : RLS(delta_, lambda_) {
        defaultParamsVector_ = initParam;
    }

    /**
     * @brief 重置参数
     *
     */
    void reset() {
        transMatrix_ = Matrix<Dim, Dim>::eye() * delta_;
        gainVector_ = Matrix<Dim, 1>::zeros();
        paramsVector_ = Matrix<Dim, 1>::zeros();
    }

    /**
     * @brief 更新参数
     *
     * @param sample 采样向量
     * @param output 实际输出
     * @return Matrix<dim, 1> const& 模型参数向量
     */
    Matrix<Dim, 1> const& update(Matrix<Dim, 1>& sample, float output) {
        gainVector_ = (transMatrix_ * sample) / (1.0F + (sample.T() * transMatrix_ * sample)[0][0] / lambda_) /
                      lambda_; ///> 计算增益向量
        paramsVector_ += gainVector_ * (output - (sample.T() * paramsVector_)[0][0]); ///> 更新参数向量
        transMatrix_ = (transMatrix_ - gainVector_ * sample.T() * transMatrix_) / lambda_; ///> 更新协方差矩阵
        return paramsVector_;
    }

    /**
     * @brief 获取参数向量
     *
     * @return Matrix<dim, 1> const& 参数向量
     */
    Matrix<Dim, 1> const& getParamVector() const { return paramsVector_; }

    /**
     * @brief 设置参数向量
     *
     * @param param 参数向量
     */
    void setParamVector(Matrix<Dim, 1> param) { paramsVector_ = param; }

    /**
     * @brief 计算输出
     *
     * @param sample 采样向量
     * @return float 当前模型下计算得到的输出
     */
    float calOutput(Matrix<Dim, 1>& sample) { return (sample.T() * paramsVector_)[0][0]; }
};
} // namespace os::math